package org.jflame.commons.crypto;

import org.jflame.commons.codec.Transcoder;
import org.jflame.commons.crypto.digest.SignatureAlg;
import org.jflame.commons.file.FileHelper;
import org.jflame.commons.util.CharsetHelper;
import org.jflame.commons.util.IOHelper;
import org.jflame.commons.util.StringHelper;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.UnrecoverableEntryException;
import java.security.UnrecoverableKeyException;
import java.security.cert.CertificateException;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Objects;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.DESKeySpec;
import javax.crypto.spec.DESedeKeySpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

public abstract class KeyUtils {

    /**
     * 生成一个密钥对
     * <p>
     * See the KeyPairGenerator section in the
     * <a href= "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyPairGenerator">
     * Java Cryptography Architecture Standard Algorithm Name Documentation</a>
     * 
     * @param keyAlgorithm 密钥算法名 DiffieHellman(DH)/DSA/RSA/EC/RSASSA-PSS
     * @param params
     * @param provider
     * @return
     * @throws EncryptException
     */
    public static KeyPair generateKeyPair(String keyAlgorithm, AlgorithmParameterSpec params, String provider)
            throws EncryptException {
        KeyPairGenerator kpg = null;
        if ("BC".equals(provider)) {
            BC.addProvider();
        }
        try {
            kpg = provider == null ? KeyPairGenerator.getInstance(keyAlgorithm)
                    : KeyPairGenerator.getInstance(keyAlgorithm, provider);
            kpg.initialize(params, new SecureRandom());
            KeyPair keyPair = kpg.generateKeyPair();
            return keyPair;
        } catch (NoSuchProviderException | NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
            throw new EncryptException(e);
        }
    }

    /**
     * 生成一个密钥对
     * 
     * @param keyAlgorithm 密钥算法名 DiffieHellman(DH)/DSA/RSA/EC/RSASSA-PSS
     * @param keysize 指定密钥大小
     * @return
     * @throws EncryptException
     */
    public static KeyPair generateKeyPair(String keyAlgorithm, int keysize, String provider) throws EncryptException {
        KeyPairGenerator kpg = null;
        if ("BC".equals(provider)) {
            BC.addProvider();
        }
        try {
            kpg = provider == null ? KeyPairGenerator.getInstance(keyAlgorithm)
                    : KeyPairGenerator.getInstance(keyAlgorithm, provider);

            kpg.initialize(keysize);
            KeyPair keyPair = kpg.generateKeyPair();
            return keyPair;
        } catch (NoSuchProviderException | NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        }
    }

    /**
     * 生成一个rsa密钥对
     * 
     * @param keysize 密钥大小,null默认2048
     * @return
     * @throws EncryptException
     */
    public static KeyPair generateRSAKeyPair(int keysize) throws EncryptException {
        return generateKeyPair(Alg.RSA.keyAlgorithm(), keysize, null);
    }

    /**
     * SM2算法生成密钥对
     *
     * @return 密钥对信息
     * @throws EncryptException
     */
    public static KeyPair generateSM2KeyPair() throws EncryptException {
        return KeyUtils.generateKeyPair("EC", new ECGenParameterSpec("sm2p256v1"), "BC");
    }

    /**
     * 从文件中加载一个Key
     * 
     * @param keyFile 文件路径,classpath:开头表示从类路径读取
     * @return
     * @throws EncryptException
     */
    public static Key loadKeyFromFile(String keyFile) throws EncryptException {
        InputStream instream = null;
        try {
            if (keyFile.startsWith("classpath:")) {
                instream = FileHelper.readFileFromClassPath(keyFile.substring(10));
            } else {
                instream = Files.newInputStream(Paths.get(keyFile));
            }
        } catch (IOException e) {
            throw new EncryptException("密钥文件加载异常", e);
        }
        return loadKey(instream);
    }

    /**
     * 从输入流生成密钥文件
     * 
     * @param keyFileStream 密钥文件流
     * @return
     * @throws EncryptException
     */
    public static Key loadKey(InputStream keyFileStream) throws EncryptException {
        ObjectInputStream ois = null;
        try {
            ois = new ObjectInputStream(keyFileStream);
            return (Key) ois.readObject();
        } catch (IOException | ClassNotFoundException e) {
            throw new EncryptException("密钥文件加载异常", e);
        } finally {
            IOHelper.closeQuietly(keyFileStream);
            IOHelper.closeQuietly(ois);
        }
    }

    /**
     * 将一个密钥写入文件中
     * 
     * @param key 密钥
     * @param keyPath 要写入的文件路径
     * @throws IOException
     */
    public static void writeKey(Key key, Path keyPath) throws IOException {
        try (ObjectOutputStream keyOutStream = new ObjectOutputStream(Files.newOutputStream(keyPath))) {
            keyOutStream.writeObject(key);
        }
    }

    /**
     * 将公私钥分别写入对应文件中
     * 
     * @param keyPair 密钥对
     * @param publicKeyPath 公钥文件路径
     * @param privateKeyPath 私钥文件路径
     * @throws IOException
     */
    public static void writeKey(KeyPair keyPair, Path publicKeyPath, Path privateKeyPath) throws IOException {
        Key publicKey = keyPair.getPublic();
        Key privateKey = keyPair.getPrivate();
        // 用对象流将生成的密钥写入文件
        try (ObjectOutputStream pubKeyOutStream = new ObjectOutputStream(Files.newOutputStream(publicKeyPath));
                ObjectOutputStream priKeyOutStream = new ObjectOutputStream(Files.newOutputStream(privateKeyPath))) {
            pubKeyOutStream.writeObject(publicKey);
            priKeyOutStream.writeObject(privateKey);
        }
    }

    /**
     * 从一个公钥文件中生成公钥对象
     * 
     * @param keyFile 公钥文件,必须是
     * @return
     * @see #writeKey(Key, Path)
     */
    public static PublicKey loadPublicKeyFromFile(String keyFile) throws EncryptException {
        return (PublicKey) loadKeyFromFile(keyFile);
    }

    /**
     * 生成公钥对象PublicKey
     * 
     * @param KeyAlgorithm 密钥算法名,如:RSA/DSA/EC
     * @param keyBytes 公钥bytes
     * @return PublicKey
     */
    public static PublicKey loadPublicKey(String KeyAlgorithm, byte[] keyBytes) throws EncryptException {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(KeyAlgorithm);
            X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
            return keyFactory.generatePublic(x509KeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new EncryptException("公钥生成异常", e);
        }
    }

    /**
     * 从公钥文件或公钥base64串中生成公钥对象
     * 
     * @param keyAlgorithm KeyPair算法名
     * @param publicKey 可以是公钥文件路径(以file:或classpath:开头)或者base64编码的公钥字符串
     * @return
     */
    public static PublicKey loadPublicKey(String keyAlgorithm, String publicKey) throws EncryptException {
        String keyPath = extractPath(publicKey);
        if (keyPath != null) {
            return loadPublicKeyFromFile(keyPath);
        }
        return loadPublicKey(keyAlgorithm, Transcoder.decodeBase64(publicKey));
    }

    private static String extractPath(String str) throws EncryptException {
        String[] prefixs = { "file:","classpath:" };
        if (str.startsWith(prefixs[0])) {
            return str.substring(prefixs[0].length());
        } else if (str.startsWith(prefixs[1])) {
            return str;
        }
        return null;
    }

    /**
     * 从公钥文件或公钥base64串中生成公钥对象
     * 
     * @param algorithm 非对称加密算法
     * @param publicKey 可以是公钥文件路径(以file:或classpath:开头)或者base64编码的公钥字符串
     * @return
     * @throws EncryptException
     */
    public static PublicKey loadPublicKey(Alg algorithm, String publicKey) throws EncryptException {
        if (algorithm.isSymmetric()) {
            throw new EncryptException(algorithm + "不是非对称算法");
        }
        return loadPublicKey(algorithm.keyAlgorithm(), publicKey);
    }

    /**
     * 从一个公钥文件中生成公钥对象
     * 
     * @param keyFile 公钥文件,必须是
     * @return
     * @see #writeKey(Key, Path)
     */
    public static PrivateKey loadPrivateKeyFromFile(String keyFile) throws EncryptException {
        return (PrivateKey) loadKeyFromFile(keyFile);
    }

    /**
     * 生成私钥对象
     * 
     * @param keyAlgorithm 密钥算法名,如:RSA/DSA/EC
     * @param keyBytes 私钥byte[]
     * @return PrivateKey
     */
    public static PrivateKey loadPrivateKey(String keyAlgorithm, byte[] keyBytes) throws EncryptException {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(keyAlgorithm);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            return keyFactory.generatePrivate(keySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new EncryptException("私钥生成异常", e);
        }
    }

    /**
     * 从私钥文件或私钥base64串中生成私钥对象
     * 
     * @param algorithm 非对称加密算法
     * @param privateKey 可以是私钥文件路径(以file:或classpath:开头)或者base64编码的私钥字符串
     * @return
     * @throws EncryptException
     */
    public static PrivateKey loadPrivateKey(Alg algorithm, String privateKey) throws EncryptException {
        return loadPrivateKey(algorithm.keyAlgorithm(), privateKey);
    }

    /**
     * 从私钥文件或私钥base64串中生成私钥对象
     * 
     * @param keyAlgorithm 密钥算法名,如:RSA/DSA/EC
     * @param privateKey 可以是私钥文件路径(以file:或classpath:开头)或者base64编码的私钥字符串
     * @return
     * @throws EncryptException
     */
    public static PrivateKey loadPrivateKey(String keyAlgorithm, String privateKey) throws EncryptException {
        String keyPath = extractPath(privateKey);
        if (keyPath != null) {
            return loadPrivateKeyFromFile(keyPath);
        }
        return loadPrivateKey(keyAlgorithm, Transcoder.decodeBase64(privateKey));
    }

    /**
     * 从公钥文件或公钥base64串中生成RSA公钥对象
     * 
     * @param publicKey 可以是公钥文件路径(以file:或classpath:开头)或者base64编码的公钥字符串
     * @return
     */
    public static PublicKey loadRSAPublicKey(String publicKey) throws EncryptException {
        return loadPublicKey(Alg.RSA.name(), publicKey);
    }

    /**
     * 从私钥文件或私钥base64串中生成RSA私钥对象
     * 
     * @param privateKey 可以是私钥文件路径(以file:或classpath:开头)或者base64编码的私钥字符串
     * @return
     * @throws EncryptException
     */
    public static PrivateKey loadRSAPrivateKey(String privateKey) throws EncryptException {
        return loadPrivateKey(Alg.RSA.name(), privateKey);
    }

    /**
     * 从公私钥文件中生成公私钥对KeyPair
     * 
     * @param publicKeyPath 公钥文件路径
     * @param privateKeyPath 私钥文件路径
     * @return KeyPair
     */
    public static KeyPair loadKeyPair(Path publicKeyPath, Path privateKeyPath) throws EncryptException {
        PublicKey publicKey = loadPublicKeyFromFile(publicKeyPath.toString());
        PrivateKey privateKey = loadPrivateKeyFromFile(privateKeyPath.toString());
        return new KeyPair(publicKey, privateKey);
    }

    public static KeyPair loadKeyPair(String publicKeyPath, String privateKeyPath) throws EncryptException {
        PublicKey publicKey = loadPublicKeyFromFile(publicKeyPath);
        PrivateKey privateKey = loadPrivateKeyFromFile(privateKeyPath);
        return new KeyPair(publicKey, privateKey);
    }

    /**
     * 根据算法自动生成一个密钥
     * 
     * @param algorithm 对称加密算法名或签名算法名
     * @return
     * @throws EncryptException
     * @see Alg
     * @see SignatureAlg
     */
    public static SecretKey generate(String algorithm) throws EncryptException {
        KeyGenerator kgen;
        try {
            if ("SM4".equals(algorithm)) {
                kgen = KeyGenerator.getInstance(algorithm, BC.addProvider());
            } else {
                kgen = KeyGenerator.getInstance(algorithm);
            }
        } catch (NoSuchAlgorithmException e) {
            throw new EncryptException(e);
        }
        return kgen.generateKey();
    }

    /**
     * 生成一个密钥
     * 
     * @param alg 对称加密算法
     * @param keyStr 密钥字符串
     * @return
     * @throws EncryptException
     */
    public static SecretKey generate(Alg alg, String keyStr) throws EncryptException {
        if (alg.name()
                .startsWith("PBE")) {
            try {
                KeySpec keySpec = new PBEKeySpec(keyStr.toCharArray());
                SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(alg.keyAlgorithm());
                return keyFactory.generateSecret(keySpec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
                throw new EncryptException(e);
            }
        } else {
            return generate(alg, CharsetHelper.getUtf8Bytes(keyStr));
        }
    }

    public static SecretKey generate(Alg alg, byte[] rawKey) throws EncryptException {
        KeySpec keySpec = null;
        try {
            if (alg == Alg.DES) {
                keySpec = new DESKeySpec(rawKey);
            } else if (alg == Alg.DESede) {
                keySpec = new DESedeKeySpec(rawKey);
            } else if (alg.name()
                    .startsWith("PBE")) {
                keySpec = new PBEKeySpec(toChars(rawKey));
            } else {
                return new SecretKeySpec(rawKey, alg.name());
            }

            SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(alg.keyAlgorithm());
            return keyFactory.generateSecret(keySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException | InvalidKeyException e) {
            throw new EncryptException(e);
        }

    }

    /**
     * 从keystore文件中读取密钥
     * 
     * @param type keystore类型
     * @param keyStoreFile keystore文件
     * @param password keystore文件保护密码
     * @param alias 密钥存储时指定的别名
     * @return 密钥 (可能是SecretKey或PrivateKey)
     * @throws KeyStoreException
     */
    public static Key getKeyFromKeystore(String type, File keyStoreFile, String password, String alias)
            throws KeyStoreException {
        try {
            KeyStore ks = readKeyStore(type, keyStoreFile, password);
            return ks.getKey(alias, password.toCharArray());
        } catch (UnrecoverableKeyException | NoSuchAlgorithmException e) {
            throw new KeyStoreException(e);
        }
    }

    public static SecretKey getSecretKeyFromKeystore(String type, File keyStoreFile, String password, String alias)
            throws KeyStoreException {
        return (SecretKey) getKeyFromKeystore(type, keyStoreFile, password, alias);
    }

    /**
     * 从keystore密钥库中读取公私钥
     * 
     * @param type
     * @param keystoreInputStream
     * @param password
     * @param alias
     * @return 公私钥KeyPair
     * @throws KeyStoreException
     */
    public static KeyPair getKeyPairFromKeystore(String type, InputStream keystoreInputStream, String password,
            String alias) throws KeyStoreException {
        char[] pwds = password.toCharArray();
        KeyStore ks = readKeyStore(type, keystoreInputStream, pwds);
        KeyStore.ProtectionParameter protectionParameter = new KeyStore.PasswordProtection(pwds);
        KeyStore.PrivateKeyEntry privateKeyEntry;
        try {
            privateKeyEntry = (KeyStore.PrivateKeyEntry) ks.getEntry(alias, protectionParameter);
            PrivateKey myPrivateKey = privateKeyEntry.getPrivateKey();
            // PrivateKey myPrivateKey2 = (PrivateKey) ks.getKey("alias", "keypass".toCharArray());
            PublicKey myPublicKey = privateKeyEntry.getCertificate()
                    .getPublicKey();
            // PublicKey myPublicKey2 = ks.getCertificate("alias").getPublicKey();
            return new KeyPair(myPublicKey, myPrivateKey);
        } catch (NoSuchAlgorithmException | UnrecoverableEntryException e) {
            throw new KeyStoreException(e);
        }
    }

    /**
     * 将SecretKey存储到keystore并保存到文件,使用pkcs12格式
     * 
     * @param key SecretKey要保存的密钥
     * @param keyStoreFile 要保存的文件
     * @param password 文件保护密码,可以是空字符但不能为null
     * @param keyAlias 密钥在keystore中的别名,用于标识该密钥
     * @throws KeyStoreException
     */
    public static void saveToKeystore(final SecretKey key, File keyStoreFile, String password, String keyAlias)
            throws KeyStoreException {
        saveToKeystore(key, "pkcs12", keyStoreFile, password, keyAlias);
    }

    /**
     * 将SecretKey存储到keystore并保存到文件
     * 
     * @param key SecretKey要保存的密钥
     * @param type keystore类型, 如果传null会使用系统默认类型,官方支持类型:jks/jceks/dks/pkcs12/pkcs11
     * @param keyStoreFile 要保存的文件
     * @param password 文件保护密码,可以是空字符但不能为null
     * @param keyAlias 密钥在keystore中的别名,用于标识该密钥
     * @throws KeyStoreException
     */
    public static void saveToKeystore(final SecretKey key, String type, File keyStoreFile, String password,
            String keyAlias) throws KeyStoreException {
        Objects.requireNonNull(password, "password can be empty but not null");
        Objects.requireNonNull(keyAlias, "keyAlias must be not null");
        char[] pwds = password.toCharArray();
        KeyStore.SecretKeyEntry secret = new KeyStore.SecretKeyEntry(key);
        KeyStore.ProtectionParameter p = new KeyStore.PasswordProtection(pwds);
        try (FileOutputStream inputStream = new FileOutputStream(keyStoreFile)) {
            KeyStore ks = KeyStore.getInstance(StringHelper.isEmpty(type) ? KeyStore.getDefaultType() : type);
            ks.load(null, null);
            ks.setEntry(keyAlias, secret, p);
            ks.store(inputStream, pwds);
        } catch (NoSuchAlgorithmException | CertificateException | IOException e) {
            throw new KeyStoreException(e);
        }
    }

    /**
     * 从一个密钥库文件中生成KeyStore对象
     * 
     * @param type KeyStore类型,官方支持类型:jks/jceks/dks/pkcs12/pkcs11
     * @param keyStoreFile 密钥库文件
     * @param password 文件保护密码
     * @return
     * @throws KeyStoreException
     * @throws IOException 密钥库文件读取失败
     */
    public static KeyStore readKeyStore(String type, File keyStoreFile, String password) throws KeyStoreException {
        try {
            return readKeyStore(type, new FileInputStream(keyStoreFile), password.toCharArray());
        } catch (FileNotFoundException e) {
            throw new KeyStoreException(e);
        }
    }

    /**
     * 从输入流中生成KeyStore对象.
     * <p>
     * {@link <a href=
     * "https://docs.oracle.com/javase/8/docs/technotes/guides/security/StandardNames.html#KeyStore">KeyStore type</a> }
     * 
     * @param type KeyStore类型, 官方支持类型:jks/jceks/dks/pkcs12/pkcs11
     * @param inputStream 密钥库输入流,读取完成后会主动关闭
     * @param password 密码
     * @return
     * @throws EncryptException 异常
     */
    public static KeyStore readKeyStore(String type, InputStream inputStream, char[] password)
            throws KeyStoreException {
        KeyStore ks;
        try {
            if ("bks".equals(type)) {
                ks = KeyStore.getInstance(type, BC.addProvider());
            } else {
                ks = KeyStore.getInstance(type);
            }
            ks.load(inputStream, password);
            return ks;
        } catch (NoSuchAlgorithmException | CertificateException | IOException e) {
            throw new KeyStoreException("读取证书异常", e);
        } catch (EncryptException e) {
            throw new KeyStoreException(e.getMessage());
        } finally {
            IOHelper.closeQuietly(inputStream);
        }
    }

    public static char[] toChars(byte[] bytes) {
        ByteBuffer bb = ByteBuffer.allocate(bytes.length);
        bb.put(bytes)
                .flip();
        CharBuffer cb = StandardCharsets.UTF_8.decode(bb);
        return cb.array();
    }

    /**
     * 将密钥对象编码为base64字符串
     * 
     * @param key
     * @return
     */
    public static String encodeBase64(Key key) {
        return Transcoder.encodeBase64URLSafeString(key.getEncoded());
    }

    /**
     * 随机生成一个指定长度的盐
     * 
     * @param saltSize
     * @return
     */
    public static byte[] genSalt(int saltSize) {
        byte[] salt = new byte[saltSize];
        try {
            SecureRandom sr = SecureRandom.getInstance("SHA1PRNG");
            sr.nextBytes(salt);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        return salt;
    }

}
